package cn.fenqing.druid.filter;

import cn.fenqing.druid.visitor.InSQLASTVisitor;
import com.alibaba.druid.filter.FilterChain;
import com.alibaba.druid.filter.FilterEventAdapter;
import com.alibaba.druid.proxy.jdbc.PreparedStatementProxy;
import com.alibaba.druid.proxy.jdbc.ResultSetProxy;
import com.alibaba.druid.proxy.jdbc.StatementProxy;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLInListExpr;
import com.alibaba.druid.sql.ast.statement.SQLSelect;
import com.alibaba.druid.sql.ast.statement.SQLSelectQuery;
import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlStatement;
import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser;
import com.alibaba.druid.sql.parser.SQLParserFeature;
import com.alibaba.druid.sql.parser.SQLStatementParser;
import com.alibaba.druid.wall.WallContext;
import com.alibaba.druid.wall.WallFilter;
import com.alibaba.druid.wall.WallProvider;

import java.sql.SQLException;
import java.util.List;

/**
 * @author fenqing
 */
public class SqlValidationFilter extends WallFilter {

    public SqlValidationFilter() {
        super();
    }

    @Override
    public ResultSetProxy preparedStatement_executeQuery(FilterChain chain, PreparedStatementProxy statement) throws SQLException {
        check(statement.getSql());
        return super.preparedStatement_executeQuery(chain, statement);
    }

    @Override
    public boolean preparedStatement_execute(FilterChain chain, PreparedStatementProxy statement) throws SQLException {
        check(statement.getSql());
        return super.preparedStatement_execute(chain, statement);
    }

    @Override
    public int preparedStatement_executeUpdate(FilterChain chain, PreparedStatementProxy statement) throws SQLException {
        check(statement.getSql());
        return super.preparedStatement_executeUpdate(chain, statement);
    }

    @Override
    public String check(String sql) throws SQLException {
        validate(sql);
        return super.check(sql);
    }


    // 校验sql
    private void validate(String sql) throws SQLException {
        SQLStatementParser parser = new MySqlStatementParser(
                sql,
                SQLParserFeature.EnableSQLBinaryOpExprGroup,
                SQLParserFeature.StrictForWall
        );
        List<SQLStatement> sqlStatements = parser.parseStatementList();
        for (SQLStatement sqlStatement : sqlStatements) {
            sqlStatement.accept(new InSQLASTVisitor());
        }
    }
}